import pandas as pd
import re


def process_school_data(input_file, output_file):
    # 加载数据并处理结构问题
    df = pd.read_csv(input_file)

    # 检测并删除重复的表头行
    header_rows = df[df.apply(lambda x: x.astype(str).str.contains('学校名称').any(), axis=1)]
    if not header_rows.empty:
        df = df.drop(header_rows.index)

    # 删除空白行
    df = df.dropna(how='all')

    # 重置索引
    df = df.reset_index(drop=True)

    # 标准化列名
    df = df.rename(columns={
        '序号': 'serial_no',
        '学校名称': 'school_name',
        '学校': 'school_name',
        '学校性质': 'school_type',
        '教职工人数': 'total_staff_count',
        '专职教师人数': 'full_time_teacher_count',
        '师生比': 'student_faculty_ratio_raw'
    })

    # 筛选公办学校
    df = df[df['school_type'] == '公办']

    # 处理师生比列
    # 保留标准冒号格式的数据
    df = df[df['student_faculty_ratio_raw'].astype(str).str.contains(r'1:\d+\.?\d*$')]

    # 提取数值部分
    df['student_faculty_ratio'] = df['student_faculty_ratio_raw'].apply(
        lambda x: float(re.search(r'1:(\d+\.?\d*)', str(x)).group(1))
    )

    # 计算估算学生总数
    df['estimated_student_count'] = df['total_staff_count'] * df['student_faculty_ratio']

    # 确保数值类型
    numeric_columns = ['total_staff_count', 'full_time_teacher_count', 'estimated_student_count']
    for col in numeric_columns:
        df[col] = pd.to_numeric(df[col], errors='coerce')

    # 处理缺失数据
    df = df.dropna(subset=numeric_columns)

    # 选择并排序最终列
    final_columns = ['serial_no', 'school_name', 'total_staff_count',
                     'full_time_teacher_count', 'estimated_student_count']
    df = df[final_columns]

    # 保存输出
    df.to_csv(output_file, index=False)
    print(f"数据处理完成，已保存至 {output_file}")


# 示例调用
if __name__ == "__main__":
    input_file = "baoshan-schools-2024.csv"
    output_file = "clean-baoshan-schools-2024.csv"
    process_school_data(input_file, output_file)